package com.wf.classRoomTest.file;

import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;

public class ImageEncryptionExample {

    private static final String AES_KEY = "woAiZhouZhouAAAA"; // 16 bytes key for AES encryption (128 bits)

    public static void main(String[] args) {
        String folderPath = "C:\\Users\\mpk16\\Pictures\\zz - 副本"; // Replace with your folder path
        File folder = new File(folderPath);

        if (!folder.isDirectory()) {
            System.out.println("Invalid directory: " + folderPath);
            return;
        }

        // Encrypt images in the folder
       encryptImages(folder);

        // Decrypt images in the folder
      //decryptImages(folder);
    }

    public static void encryptImages(File folder) {
        File[] imageFiles = folder.listFiles((dir, name) -> name.toLowerCase().endsWith(".jpg")
                                                        || name.toLowerCase().endsWith(".jpeg")
                                                        || name.toLowerCase().endsWith(".png"));

        if (imageFiles == null) {
            System.out.println("No image files found in: " + folder.getAbsolutePath());
            return;
        }

        for (File imageFile : imageFiles) {
            try {
                byte[] imageData = Files.readAllBytes(imageFile.toPath());
                byte[] encryptedData = encryptData(imageData, AES_KEY.getBytes());

                // Write encrypted data back to the image file
                try (FileOutputStream fos = new FileOutputStream(imageFile)) {
                    fos.write(encryptedData);
                }

                System.out.println("Encrypted: " + imageFile.getName());
            } catch (IOException | NoSuchAlgorithmException
                     | NoSuchPaddingException | InvalidKeyException
                     | IllegalBlockSizeException | BadPaddingException e) {
                e.printStackTrace();
            }
        }
    }

    public static void decryptImages(File folder) {
        File[] imageFiles = folder.listFiles((dir, name) -> name.toLowerCase().endsWith(".jpg")
                                                        || name.toLowerCase().endsWith(".jpeg")
                                                        || name.toLowerCase().endsWith(".png"));

        if (imageFiles == null) {
            System.out.println("No image files found in: " + folder.getAbsolutePath());
            return;
        }

        for (File imageFile : imageFiles) {
            try {
                byte[] encryptedData = Files.readAllBytes(imageFile.toPath());
                byte[] decryptedData = decryptData(encryptedData, AES_KEY.getBytes());

                // Write decrypted data back to the image file
                try (FileOutputStream fos = new FileOutputStream(imageFile)) {
                    fos.write(decryptedData);
                }

                System.out.println("Decrypted: " + imageFile.getName());
            } catch (IOException | NoSuchAlgorithmException
                    | NoSuchPaddingException | InvalidKeyException
                    | IllegalBlockSizeException | BadPaddingException e) {
                e.printStackTrace();
            }
        }
    }

    public static byte[] encryptData(byte[] data, byte[] key)
            throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException,
            IllegalBlockSizeException, BadPaddingException {
        SecretKeySpec secretKeySpec = new SecretKeySpec(key, "AES");
        Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
        return cipher.doFinal(data);
    }

    public static byte[] decryptData(byte[] encryptedData, byte[] key)
            throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException,
            IllegalBlockSizeException, BadPaddingException {
        SecretKeySpec secretKeySpec = new SecretKeySpec(key, "AES");
        Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
        return cipher.doFinal(encryptedData);
    }
}
